package com.yupi.yuaiagent.demo.rag;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.preretrieval.query.transformation.QueryTransformer;
import org.springframework.ai.rag.preretrieval.query.transformation.RewriteQueryTransformer;
import org.springframework.stereotype.Component;

/**
 * 查询重写器
 * @className: QueryRewriter
 * @author: xxy-Rain
 * @date: 2025/10/24 10:51
 * @version: 1.0
 * @description: TODO
 */
@Component
public class QueryRewriter {
    private final QueryTransformer queryTransformer;
    public QueryRewriter(ChatModel dashscopeChatModel){
        ChatClient.Builder builder = ChatClient.builder(dashscopeChatModel);
        //创建查询重写转换器
        queryTransformer = RewriteQueryTransformer.builder()
                .chatClientBuilder(builder)
                .build();
    }

    /**
     * 执行查询重写
     * @param prompt
     * @return
     */
    public String doQueryRewrite(String prompt){
        Query query = new Query(prompt);
        //查询重写
        Query transformedQuery = queryTransformer.transform(query);
        //返回重写后的查询
        return transformedQuery.text();
    }
}
